package uk.ac.ebi.pride.peptidome;

import org.apache.log4j.Logger;
import uk.ac.ebi.pride.jaxb.model.*;
import uk.ac.ebi.pride.jaxb.xml.marshaller.PrideXmlMarshaller;
import uk.ac.ebi.pride.jaxb.xml.marshaller.PrideXmlMarshallerFactory;
import uk.ac.ebi.pride.peptidome.model.MSProtein;
import uk.ac.ebi.pride.tools.mgf_parser.MgfFile;
import uk.ac.ebi.pride.tools.mgf_parser.model.Ms2Query;
import uk.ac.ebi.pride.tools.peak_list_parser.PeakListParserException;

import java.io.*;
import java.lang.management.MemoryUsage;
import java.util.*;
import java.util.zip.GZIPOutputStream;

/**
 * Created by IntelliJ IDEA.
 * User: rwang
 * Date: 20/05/11
 * Time: 10:47
 * To change this template use File | Settings | File Templates.
 */
public class PeptidomePrideXmlWriter {
    private static final Logger logger = Logger.getLogger(PeptidomePrideXmlWriter.class);

    private final String outputFilePath;

    private final String prideAccession;

    private final PrideXmlMarshaller marshaller;

    private final PeptidomeToPrideInterpreter interpreter;

    private final Map<String, MgfFile> mgfFiles;

    private final PeptidomeReader peptidomeReader;

    private final String sampleAcc;

    public PeptidomePrideXmlWriter(PeptidomeToPrideInterpreter interpreter, PeptidomeReader peptidomeReader, String sampleAcc,
                                   Map<String, MgfFile> mgfFiles, String outputFilePath, String prideAcc) {
        // output file path
        this.outputFilePath = outputFilePath;
        // pride xml marshaller
        this.marshaller = PrideXmlMarshallerFactory.getInstance().initializeMarshaller();
        // peptidome to pride interpreter
        this.interpreter = interpreter;
        // a map of mgf file names and mgf file readers
        this.mgfFiles = mgfFiles;
        // reader for getting information from peptidome file
        this.peptidomeReader = peptidomeReader;
        // the accession of the sample
        this.sampleAcc = sampleAcc;
        // pride accession
        this.prideAccession = prideAcc;
    }

    public void writerXml() throws ConverterException {
        PrintWriter out = null;
        try {

            //create xml file
            File xmlFile = new File(outputFilePath);

            out = new PrintWriter(new FileWriter(xmlFile));

            //write header
            out.println("<?xml version=\"1.0\" encoding=\"UTF-8\"?>");

            //experiment
            out.println("<ExperimentCollection version=\"2.1\">");
            out.println("<Experiment>");
            // experiment accession
            out.println("<ExperimentAccession>" + prideAccession + "</ExperimentAccession>");
            // experiment title
            out.println("<Title>" + interpreter.getExperimentTitle() + "</Title>");
            logger.info("Finish writing experiment title");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));            // experiment reference
            marshallPrideObject(out, interpreter.getReferences());
            logger.info("Finish writing references");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));            // experiment short label
            out.println("<ShortLabel>" + interpreter.getSampleAccession() + "</ShortLabel>");
            logger.info("Finish writing short label");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));            //protocol
            marshallPrideObject(out, interpreter.getProtocol());
            out.println();
            logger.info("Finish writing protocol");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));
            //mzdata
            out.println("<mzData version=\"1.05\" accessionNumber=\"\">");
            marshallPrideObject(out, interpreter.getCvLookup());
            out.println();
            out.println("<description>");
            marshallPrideObject(out, interpreter.getAdmin());
            out.println();
            marshallPrideObject(out, interpreter.getInstrument());
            out.println();
            marshallPrideObject(out, interpreter.getDataProcessing());
            out.println();
            out.println("</description>");
            logger.info("Finish writing mzData description");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));
            //spectrumlist - note dependency on includeAllSpectra
            out.println("<spectrumList count=\"" + interpreter.getSpectrumCount(mgfFiles.values()) + "\">");
            Map<String, Integer> spectrumIdMap = new HashMap<String, Integer>();
            int spectrumCount = 0;
            Iterator<Map.Entry<String, MgfFile>> spectrumIter = mgfFiles.entrySet().iterator();
            while (spectrumIter.hasNext()) {
                Map.Entry<String, MgfFile> entry = spectrumIter.next();
                String mgfFileName = entry.getKey();
                MgfFile mgfFile = entry.getValue();
                Iterator<Ms2Query> ms2QueryIterator = null;
                try {
                    ms2QueryIterator = mgfFile.getMs2QueryIterator();
                } catch (PeakListParserException e) {
                    logger.error("Failed to get spectrum iterator", e);
                }
                while (ms2QueryIterator.hasNext()) {
                    Ms2Query ms2 = ms2QueryIterator.next();
                    Spectrum spectrum = interpreter.createSpectrumFromMs2Query(ms2);
                    String peptidomeSpectrumId = mgfFileName + ":" + ms2.getTitle().replace("Scan", "").trim();
                    int spectrumId;
                    if (spectrumIdMap.containsKey(peptidomeSpectrumId)) {
                        spectrumId = spectrumIdMap.get(peptidomeSpectrumId);
                    } else {
                        spectrumCount++;
                        spectrumIdMap.put(peptidomeSpectrumId, spectrumCount);
                        spectrumId = spectrumCount;
                    }
                    spectrum.setId(spectrumId);
                    spectrum.getSpectrumDesc().getComments().add(peptidomeSpectrumId);
                    marshallPrideObject(out, spectrum);
                    out.println();
                    out.flush();
                }

                spectrumIter.remove();
            }
            //close spectrumlist
            out.println("</spectrumList>");
            logger.info("Finish writing spectrum list");
            logger.info("Used memory: " + (Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory()));
            //close mzdata
            out.println("</mzData>");

            //identifications - we iterate from the reader to ensure that we capture any changes made
            //                  by external tools (e.g. protein inference)

            List<String> proteinIds = peptidomeReader.getSampleProteinIds(sampleAcc);
            if (proteinIds != null) {
                for (String proteinId : proteinIds) {
                    MSProtein protein = peptidomeReader.getSampleProteinById(sampleAcc, proteinId);
                    List<Identification> identifications = interpreter.convertToPrideProteins(protein, spectrumIdMap);
                    for (Identification identification : identifications) {
                        marshallPrideObject(out, identification);
                        out.println();
                    }
                    out.flush();
                }
            }

            //close experiment
            marshallPrideObject(out, interpreter.getExperimentAdditional());
            out.println("</Experiment>");

            //close experimentcollection
            out.println("</ExperimentCollection>");

        } catch (IOException e) {
            throw new ConverterException("Error writing output file: " + e.getMessage(), e);
        } finally {
            if (out != null) {
                out.flush();
                out.close();
            }
        }
    }


    /**
     * helper method to marshall out a collection of PRIDE JAXB model objects
     *
     * @param out
     * @param xmlObject
     * @param <T>
     * @throws IOException
     */
    private <T extends PrideXmlObject> void marshallPrideObject(PrintWriter out, List<T> xmlObject) throws IOException {
        if (xmlObject != null) {

            //iterate and marshall individual objects
            for (T obj : xmlObject) {
                marshallPrideObject(out, obj);
                out.println();
            }
        }

    }

    /**
     * helper method to marshall out a single PRIDE JAXB model objects
     *
     * @param out
     * @param xmlObject
     * @param <T>
     * @throws IOException
     */
    private <T extends PrideXmlObject> void marshallPrideObject(PrintWriter out, T xmlObject) throws IOException {

        if (xmlObject != null) {
            try {
                marshaller.marshall(xmlObject, out);
            } catch (Exception e) {
                throw new IOException("Error marshalling: " + xmlObject.getClass().getName(), e);
            }
        }
    }
}
